import os
from typing import Optional, List

import chainlit as cl
from chainlit.element import ElementBased
from dotenv import load_dotenv

from llama_index.core.chat_engine import SimpleChatEngine
from llama_index.core.chat_engine.types import ChatMode
from rag.traditional_rag import TraditionalRAG
from persistent.minio_storage_client import MinioStorageClient
from persistent.postgresql_data_layer import PostgreSQLDataLayer
import chainlit.data as cl_data
from rag.config import RagConfig
load_dotenv()
# 实现聊天数据持久化
storage_client = MinioStorageClient()
cl_data._data_layer = PostgreSQLDataLayer(conninfo=RagConfig.pg_connection_string, storage_provider=storage_client)



async def view_pdf(elements: List[ElementBased]):
    """查看PDF文件"""
    files = []
    contents = []
    for element in elements:
        if element.name.endswith(".pdf"):
            pdf = cl.Pdf(name=element.name, display="side", path=element.path)
            files.append(pdf)
            contents.append(element.name)
    if len(files) == 0:
        return
    await cl.Message(content=f"查看PDF文件：" + "，".join(contents), elements=files).send()


@cl.on_chat_start
async def start():
    # 直接与大模型对话引擎
    chat_engine = SimpleChatEngine.from_defaults()
    cl.user_session.set("chat_engine", chat_engine)

    await cl.Message(
        author="Assistant", content="您好，我是但问智能助手，请问有什么可以帮到您的吗？"
    ).send()


@cl.on_message
async def main(message: cl.Message):
    chat_engine = cl.user_session.get("chat_engine")
    msg = cl.Message(content="", author="Assistant")
    # 预览pdf
    await view_pdf(message.elements)

    files = []
    # 获取用户上传的文件（包含图片）
    for element in message.elements:
        if isinstance(element, cl.File) or isinstance(element, cl.Image):
            files.append(element.path)
    # 文件索引处理
    if len(files) > 0:
        # 构建索引，正常执行一次即可
        rag = TraditionalRAG(files=files)
        index = await rag.create_index_local()
        chat_engine = index.as_chat_engine(chat_mode=ChatMode.CONTEXT)
        cl.user_session.set("chat_engine", chat_engine)

    res = await cl.make_async(chat_engine.stream_chat)(message.content)

    # 流式界面输出
    for token in res.response_gen:
        await msg.stream_token(token)
    await msg.send()

@cl.password_auth_callback
def auth_callback(username: str, password: str) -> Optional[cl.User]:
    # 可以对接第三方认证
    if (username, password) == ("admin", "admin"):
        return cl.User(identifier="admin",
                       metadata={"role": "admin", "provider": "credentials"})
    else:
        return None